import logging

from . import Shellcodes

l = logging.getLogger(name=__name__)


class ShellcodeFactory:
    """
    Shellcode manager, manages preferences for shellcode types and in the future, their options

    :ivar os:       Platform of this shellcode factory.
    :ivar arch:     Architecture of this shellcoe factory.
    :type arch:     archinfo.Arch
    :ivar default:  A dict of default shellcode.
    :ivar shellcode_args:   Arguments that will be passed to default shellcode when initializing.
    """

    DEFAULT_SHELLCODES = {'cgc': 'setregister', 'unix': 'binsh'}

    def __init__(self, project, default=None, shellcode_args=None):
        """
        :param project: a project option to base shellcode retrieve decisions off of
        :param default: a default shellcode to use
        :param shellcode_args: arguments to pass to the default shellcode when constructing
        """
        self.os = 'unix' if project.loader.main_object.os.startswith('UNIX') else project.loader.main_object.os
        self.arch = project.arch

        self.default = dict(ShellcodeFactory.DEFAULT_SHELLCODES)
        if not default is None:
            if self.check_shellcode_existence(default):
                self.default[self.os] = default

        self.shellcode_args = { } if shellcode_args is None else shellcode_args

    def check_shellcode_existence(self, name):
        if not name in Shellcodes[self.os][self.arch.name]:
            l.warning("shellcode by name '%s' for os '%s' and arch '%s' does not exist", name, self.os, self.arch)
            return False
        return True

    def get_default(self):
        """
        retrieve the default shellcode as a raw string
        """
        payload_name = self.default[self.os]
        sc = Shellcodes[self.os][self.arch.name][payload_name](**self.shellcode_args)
        return sc.raw(arch=self.arch)

    def set_default(self, new_default):
        """
        set a new default shellcode
        """
        self.check_shellcode_existence(new_default)
        self.default[self.os] = new_default

    def get_shellcode(self, payload_name, **kwargs):
        """
        grab a shellcode with a known name
        """
        self.check_shellcode_existence(payload_name)
        sc = Shellcodes[self.os][self.arch.name][payload_name](**kwargs)
        return sc.raw(arch=self.arch)

    def get_shellcode_template(self, payload_name, **kwargs):
        """
        grab a shellcode with a known name, do not convert it into raw code or template it
        """
        self.check_shellcode_existence(payload_name)
        return Shellcodes[self.os][self.arch.name][payload_name](**kwargs)
